#pragma once

#include <Magick++.h>

#include "LBM.h"

struct VelMagnitudePainter
{
  LBfloat m_scale;

  explicit VelMagnitudePainter(LBfloat scale = 1.0) :  m_scale(scale) {}

  Magick::Color operator() (const LatticeBoltzmannD2Q7 & ca, int x, int y) const
  {
      const D2Q7Cell & cell = ca.getCell(x, y);
      LBfloat d = calcDensity(cell);
      LBfloat ux, uy;
      calcMomentum(cell, ux, uy);
      if (d > eps)
      {
        ux /= d;
        uy /= d;
      }

      LBfloat v = sqrt(ux*ux + uy*uy) * m_scale;
      return Magick::ColorRGB(v, v, v);
  }
};

struct VelHSLPainter
{
  LBfloat m_scale;

  explicit VelHSLPainter(LBfloat scale = 1.0) :  m_scale(scale) {}

  Magick::Color operator() (const LatticeBoltzmannD2Q7 & ca, int x, int y) const
  {
    if (ca.getCellType(x, y) == CELL_WALL)
      return Magick::ColorRGB(0.5, 0.5, 0.5);

    const D2Q7Cell & cell = ca.getCell(x, y);
    LBfloat d = calcDensity(cell);
    LBfloat ux, uy;
    calcMomentum(cell, ux, uy);
    if (d > eps)
    {
      ux /= d;
      uy /= d;
    }

    LBfloat ang = atan2(ux, uy) / (2.0 * PI) + 0.5;
    LBfloat v = sqrt(ux*ux + uy*uy) * m_scale;
    return Magick::ColorHSL(ang, 1.0, 0.5 * v);
  }
};

Magick::Image renderVelocityMap(int size, LBfloat m_scale = 1.0)
{
  int width = size;
  int height = size;

  Magick::Image img;
  img.size( Magick::Geometry(width, height) );
  img.modifyImage();
  Magick::PixelPacket *pixel_cache = img.getPixels(0, 0, width, height);
  
  for (int y = 0; y != height; ++y)
    for (int x = 0; x != width; ++x)
    {
      Magick::PixelPacket & p = *(pixel_cache + y * width + x);
      LBfloat ux = x / (0.5 * width) - 1.0;
      LBfloat uy = (y / (0.5 * height) - 1.0);
      LBfloat ang = atan2(ux, uy) / (2.0 * PI) + 0.5;
      LBfloat v = sqrt(ux*ux + uy*uy);
      if (v <= 1.0)
        p = Magick::ColorHSL(ang, 1.0, 0.5 * v * m_scale);
      else
        p = Magick::ColorRGB(0, 0, 0);
    }

  return img;
}
